package com.example.demo1.samplews;

import com.alibaba.fastjson.JSONObject;
import lombok.Data;
import okhttp3.Response;
import okhttp3.ResponseBody;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;

public class AiAgentHttp {

    private String url;

    private String model;

    private int maxRound;

    private String firstSentence;

    private String promptSystem;

    List<JSONObject> messages = new ArrayList<>();

    AiAgentHttp partner = null;

    public AiAgentHttp(String url, String model, int maxRound, String firstSentence, String promptSystem) {
        this.url = url;
        this.model = model;
        if (maxRound < 3) {
            throw new RuntimeException("对话轮数不能小于3");
        }
        this.maxRound = maxRound;
        this.firstSentence = firstSentence;
        this.promptSystem = promptSystem;
        JSONObject promptSystemMessage = new JSONObject();
        promptSystemMessage.put("role", "system");
        promptSystemMessage.put("content", promptSystem);
        this.messages.add(promptSystemMessage);
        System.out.println("【" + model + "】 已激活");
    }

    //开启对话
    public void start() {
        if (this.getPartner() == null) {
            throw new RuntimeException("没有可对话的aiAgent");
        }
        System.out.println("【" + this.model + "】： " + firstSentence);
        this.partner.listenThenTalk(firstSentence);
    }

    //听，说
    public void listenThenTalk(String msg) {
        //记录user的话
        JSONObject userAgentSaid = new JSONObject();
        userAgentSaid.put("role", "user");
        userAgentSaid.put("content", msg);
        if (this.messages.size() > maxRound) {
            this.messages.remove(this.messages.get(1));
        }
        this.messages.add(userAgentSaid);
        JSONObject param = new JSONObject();
        param.put("model", this.model.trim());
        param.put("messages", this.messages);
        param.put("stream", true);
        Response response = null;
        InputStream inputStream = null;
        InputStreamReader isr = null;
        BufferedReader br = null;
        try {
            //向assistant发送消息
            response = OKHttpUtil.getOutputStream(this.url, param.toString());
            //处理assistant返回的消息
            ResponseBody responseBody = response.body();
            inputStream = responseBody.byteStream();
            isr = new InputStreamReader(inputStream, Charset.defaultCharset());
            br = new BufferedReader(isr);
            System.out.print("【" + this.model + "】： ");
            String str;
            StringBuilder userAgentSaidSb = new StringBuilder();
            while ((str = br.readLine()) != null) {
                JSONObject answerObj = JSONObject.parseObject(str);
                String content = answerObj.getJSONObject("message").getString("content");
                System.out.print(content);
                userAgentSaidSb.append(content);
//                Thread.sleep(2);
            }
            System.out.println();
            //记录assistant的话
            JSONObject assistantAgentSaid = new JSONObject();
            assistantAgentSaid.put("role", "user");
            assistantAgentSaid.put("content", userAgentSaidSb.toString());
            if (this.messages.size() > maxRound) {
                this.messages.remove(this.messages.get(1));
            }
            this.messages.add(assistantAgentSaid);
            //assistant向user反问
            this.partner.listenThenTalk(userAgentSaidSb.toString());
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                inputStream.close();
                isr.close();
                br.close();
                response.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public String getUrl() {
        return url;
    }

    public void setUrl(String url) {
        this.url = url;
    }

    public String getModel() {
        return model;
    }

    public void setModel(String model) {
        this.model = model;
    }

    public int getMaxRound() {
        return maxRound;
    }

    public void setMaxRound(int maxRound) {
        this.maxRound = maxRound;
    }

    public String getFirstSentence() {
        return firstSentence;
    }

    public void setFirstSentence(String firstSentence) {
        this.firstSentence = firstSentence;
    }

    public String getPromptSystem() {
        return promptSystem;
    }

    public void setPromptSystem(String promptSystem) {
        this.promptSystem = promptSystem;
    }

    public List<JSONObject> getMessages() {
        return messages;
    }

    public void setMessages(List<JSONObject> messages) {
        this.messages = messages;
    }

    public AiAgentHttp getPartner() {
        return partner;
    }

    public void setPartner(AiAgentHttp partner) {
        this.partner = partner;
    }
}
